package main

import (
	"crypto/tls"
	"log"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type WsServer struct {
	ws_mu           sync.Mutex
	conn_list       []*websocket.Conn
	data_ch         chan int
	server          *http.Server
	consuming       bool
	sub_leader      string
	sub_leader_conn *websocket.Conn
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	// 允许跨域访问
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func (server *WsServer) remove_conn(conn *websocket.Conn) {
	server.ws_mu.Lock()
	defer server.ws_mu.Unlock()
	idx := -1
	for i := range server.conn_list {
		if server.conn_list[i] == conn {
			idx = i
			break
		}
	}

	if idx == -1 {
		return
	}

	server.conn_list = append(server.conn_list[:idx], server.conn_list[idx+1:]...)
}

func (server *WsServer) clientWs(w http.ResponseWriter, r *http.Request) {
	server.ws_mu.Lock()
	defer server.ws_mu.Unlock()
	sub_leader := r.FormValue("sub_leader")
	if server.consuming && server.sub_leader == sub_leader {
		w.Write([]byte("Already comsuming."))
		return
	}

	server.consuming = true
	server.sub_leader = sub_leader

	go server.consume_data(sub_leader)
	w.Write([]byte("Comsuming boot."))
}

func (server *WsServer) serveWs(w http.ResponseWriter, r *http.Request) {
	server.ws_mu.Lock()
	// 升级HTTP连接到WebSocket连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Failed to set websocket upgrade:", err)
		return
	}
	server.conn_list = append(server.conn_list, conn)
	server.ws_mu.Unlock()
}

func (server *WsServer) produce_data() {
	for {
		select {
		case <-time.After(2000 * time.Millisecond):
			num := rand.Intn(100000)
			server.data_ch <- num
		}
	}
}

func (server *WsServer) consume_data(sub_leader string) {
	server.ws_mu.Lock()
	if server.sub_leader_conn != nil {
		server.sub_leader_conn.Close()
	}
	server.ws_mu.Unlock()

	ws_url := "ws://" + sub_leader + "/serveWs"
	conn, _, err := websocket.DefaultDialer.Dial(ws_url, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	server.sub_leader_conn = conn

	// 读取消息
	for {
		_, p, err := conn.ReadMessage()
		if err != nil {
			log.Fatal(err)
		}
		log.Println("Received message:", string(p))
	}
}

func (server *WsServer) send_conns_data() {
	for data := range server.data_ch {
		server.ws_mu.Lock()
		for i := range server.conn_list {
			err := server.conn_list[i].WriteJSON(data)
			if err != nil {
				log.Println("Error writing message:", err)
			}
		}
		server.ws_mu.Unlock()
	}
}

func boot_websocket() {
	port := "11222"
	tlsClientSkipVerify := &tls.Config{InsecureSkipVerify: true}

	ws_server := WsServer{}
	ws_server.consuming = false
	ws_server.conn_list = make([]*websocket.Conn, 0)
	ws_server.data_ch = make(chan int, 256)
	ws_server.sub_leader = ""
	ws_server.sub_leader_conn = nil

	http.HandleFunc("/clientWs", ws_server.clientWs)
	http.HandleFunc("/serveWs", ws_server.serveWs)
	ws_server.server = &http.Server{
		Addr:      ":" + port,
		TLSConfig: tlsClientSkipVerify,
	}

	go ws_server.produce_data()
	go ws_server.send_conns_data()

	if err := ws_server.server.ListenAndServe(); err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

func main() {
	boot_websocket()
}
